# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jhkim

import collections
import functools
import os
import datetime
import logging
import sys
import numpy
import pandas

from absl import app, flags

from coin.base.book.types import TradeSide
from coin.exchange.shared.kr_rest.product import SharedProduct
from coin.exchange.okex_futures.kr_rest.futures_product import OkexFuturesProduct
from coin.exchange.bitmex.kr_rest.futures_product import BitmexFuturesProduct
from coin.exchange.bitflyer_v1.kr_rest.futures_product import BitflyerFuturesProduct
from coin.support.proto_log.logic.util import run_from_og_log_archive
from coin.strategy.mm.dumper_base import run_from_archive
from coin.strategy.accounting.fee import get_pnl_fee
from coin.proto.coin_order_gateway_pb2 import OrderGatewayLog, OrderEvent
from coin.proto.coin_order_enums_pb2 import (BUY_ORDER,
                                             SELL_ORDER,
                                             UNKNOWN_ORDER_SIDE,
                                             BUY_OPEN_ORDER,
                                             BUY_CLOSE_ORDER,
                                             SELL_OPEN_ORDER,
                                             SELL_CLOSE_ORDER)
from coin.strategy.mm.tool.duration_tracker import DurationTracker
import google.protobuf.json_format as json_format

import coin.strategy.mm.feed as feed

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

flags.DEFINE_bool('detail', False, 'show detail stats')

FLAGS = flags.FLAGS


def get_weighted_mean_robust(values, weights):
  if weights.sum() > 0:
    return numpy.average(values, weights=weights)
  else:
    return float('nan')


def convert_side(side):
  if side == 'BUY_ORDER':
    sign = 1
    side = BUY_ORDER
  elif side == 'BUY_OPEN_ORDER':
    sign = 1
    side = BUY_OPEN_ORDER
  elif side == 'BUY_CLOSE_ORDER':
    sign = 1
    side = BUY_CLOSE_ORDER
  elif side == 'SELL_ORDER':
    sign = -1
    side = SELL_ORDER
  elif side == 'SELL_OPEN_ORDER':
    sign = -1
    side = SELL_OPEN_ORDER
  elif side == 'SELL_CLOSE_ORDER':
    sign = -1
    side = SELL_CLOSE_ORDER
  elif side == 'UNKNOWN_ORDER_SIDE':
    sign = 0
    side = UNKNOWN_ORDER_SIDE
  else:
    raise ValueError(side)
  return sign, side


class OrderIdConverter(object):
  def __init__(self):
    self._internal_id_map = collections.defaultdict(dict)
    self._last_time = collections.defaultdict(int)
    self._order_price_oid_map = collections.defaultdict(int)

  def get_order_id(self, event):
    if self._last_time[event.symbol] > event.event_time:
      return None
    self._last_time[event.symbol] = event.event_time
    return self.get_order_id_impl(event, self._internal_id_map[event.symbol])

  def get_order_id_impl(self, event, internal_id_map):
    if event.type == OrderEvent.ORDER_SUBMITTED:
      if len(internal_id_map) == 0:
        order_id = 0
      else:
        order_id = max(internal_id_map.values()) + 1
      internal_id_map[event.internal_order_id] = order_id
    if event.internal_order_id not in internal_id_map:
      logging.debug("invalid internal order id: %s", event.internal_order_id)
      return None
    return internal_id_map[event.internal_order_id]

  def get_order_id_from_json(self, event_json):
    symbol = event_json['symbol']
    event_time = int(event_json['event_time'])
    if self._last_time[symbol] > event_time:
      return None
    self._last_time[symbol] = event_time
    return self.get_order_id_impl_from_json(event_json, self._internal_id_map[symbol])

  def get_order_id_impl_from_json(self, event_json, internal_id_map):
    internal_order_id = event_json['internal_order_id']
    if event_json['type'] == 'ORDER_SUBMITTED':
      if len(internal_id_map) == 0:
        order_id = 0
      else:
        order_id = max(internal_id_map.values()) + 1
      internal_id_map[internal_order_id] = order_id
    if internal_order_id not in internal_id_map:
      logging.debug("invalid internal order id: %s", internal_order_id)
      return None
    return internal_id_map[internal_order_id]

  def get_order_id_from_pb(self, event_pb):
    symbol = event_pb.symbol
    event_time = int(event_pb.event_time)
    if self._last_time[symbol] > event_time:
      return None
    self._last_time[symbol] = event_time

    # bithumb case.
    if event_pb.type == OrderEvent.ORDER_FILLED and event_pb.external_order_id == "None":
      if (symbol, event_pb.order_price, event_pb.order_side) in self._order_price_oid_map:
        return self._order_price_oid_map[(symbol, event_pb.order_price, event_pb.order_side)]

    oid = self.get_order_id_impl_from_pb(event_pb, self._internal_id_map[symbol])

    if event_pb.type == OrderEvent.ORDER_ACCEPTED:
      self._order_price_oid_map[(symbol, event_pb.order_price, event_pb.order_side)] = oid

    return oid

  def get_order_id_impl_from_pb(self, event_pb, internal_id_map):
    internal_order_id = event_pb.internal_order_id
    if event_pb.type == OrderEvent.ORDER_SUBMITTED:
      if len(internal_id_map) == 0:
        order_id = 0
      else:
        order_id = max(internal_id_map.values()) + 1
      internal_id_map[internal_order_id] = order_id
    if internal_order_id not in internal_id_map:
      logging.debug("invalid internal order id: %s", internal_order_id)
      return None
    return internal_id_map[internal_order_id]


class PtaStrategy(object):
  def __init__(self, market_type, exchange, trading_date, machine):
    self._trading_date = trading_date
    self._market_type = market_type
    self._exchange = exchange
    self._machine = machine
    self._header = [
        "timestamp",
        "symbol",
        "event_type",
        "price",
        "qty",
        "side",
        "sign",
        "order_id",
        "fill_price",
        "fill_qty",
        "last_pos"
    ]
    self._rows = collections.defaultdict(list)
    self._symbol_count = collections.defaultdict(int)
    self._type_dict = {
        key: value for value,
        key in zip(OrderEvent.OrderEventType.keys(), OrderEvent.OrderEventType.values())
    }
    self._order_id_converter = OrderIdConverter()
    self._stat_oe = {}
    self._product_map = {}
    self._symbol_map = {}
    self._products = []
    self._symbol_mtm = {}
    self._stat_df = {}
    self._balance_map = {}
    self._currency_map = {}
    self._arraydict = collections.defaultdict(dict)
    self._first_log_timestamp = None
    self._last_log_timestamp = None
    self._duration_tracker = DurationTracker()
    self._balances = collections.defaultdict(list)
    self._is_okex_futures = exchange == 'Okex' and market_type == 'Futures'
    self._last_book_timestamp = 0
    self._last_trade_timestamp = 0

  def get_stat_oe(self, oedf, need_plot=True, is_okex_futures=False):
    assert len(oedf['symbol'].unique()) == 1
    if is_okex_futures:
      if oedf['symbol'][0].find('BTC') >= 0:
        multiplier = 100
      else:
        multiplier = 10
      for column in oedf.columns:
        if column.find("qty") >= 0 or column.find("pos") >= 0:
          oedf[column] = oedf[column] / oedf['price'] * multiplier

    statcol = {}

    if self._market_type == 'Spot':
      pos_reserve = (numpy.percentile(oedf['last_pos'], q=99)
                     + numpy.percentile(oedf['last_pos'], q=1)) * 0.5
      oedf['last_pos'] -= pos_reserve
    else:
      pos_reserve = 0

    filldf = oedf.loc[oedf['event_type'] == OrderEvent.ORDER_FILLED].reset_index(drop=True)

    rows = []

    if len(filldf) > 0:
      last_fill_price = filldf.iloc[-1]['fill_price']
      prc_diff = (filldf.iloc[-1]['fill_price'] - filldf.iloc[0]['fill_price'])
    else:
      last_fill_price = None
      prc_diff = None

    if pos_reserve != 0:
      statcol['pos_reserve'] = pos_reserve
      pnl_pos_reserve = pos_reserve * (prc_diff or 0)

    next_mtmp = None
    next_mtmps = []

    subdf = oedf.loc[oedf['event_type'] == OrderEvent.ORDER_SUBMITTED].reset_index(drop=True)
    candf = oedf.loc[oedf['event_type'] == OrderEvent.CANCEL_CONFIRMED].reset_index(drop=True)

    statcol.update({
        'symbol': oedf['symbol'].unique()[0],
        'lot_size_avg': subdf['qty'].mean(),
        'cnt_fill': filldf.shape[0],
        'cnt_order_sub': subdf.shape[0],
    })
    for _, row in filldf.iloc[::-1].iterrows():
      next_mtmps.append(row['fill_price'])
      if row['last_pos'] > 0:
        ret = 0 if next_mtmp is None else next_mtmp - row['fill_price']
      elif row['last_pos'] < 0:
        ret = 0 if next_mtmp is None else row['fill_price'] - next_mtmp
      else:
        ret = 0
        ret30 = 0
        ret60 = 0

      if row['sign'] > 0:
        ret30 = 0 if next_mtmp is None else numpy.mean(next_mtmps[-30:]) - row['fill_price']
        ret60 = 0 if next_mtmp is None else numpy.mean(next_mtmps[-60:]) - row['fill_price']
      else:
        ret30 = 0 if next_mtmp is None else row['fill_price'] - numpy.mean(next_mtmps[-30:])
        ret60 = 0 if next_mtmp is None else row['fill_price'] - numpy.mean(next_mtmps[-60:])

      row['next_fill_ret'] = ret
      row['next_fill_ret30'] = ret30
      row['next_fill_ret60'] = ret60
      row['last_fill_diff'] = last_fill_price - row['fill_price']
      next_mtmp = row['fill_price']
      rows.append(row)
    filldf = pandas.DataFrame(rows[::-1])

    try:
      first_pos = oedf['last_pos'][oedf['last_pos'].first_valid_index()]
    except Exception:
      first_pos = float('nan')

    if FLAGS.detail:
      statcol.update({
          'cancel_cnt': candf.shape[0],
          'cancel_qty': candf['qty'].sum(),
          'order_qty': subdf['qty'].sum(),
      })

    if len(filldf) > 0:
      if (len(numpy.unique(filldf['fill_qty'])) > 1
          and len(numpy.unique(filldf['next_fill_ret'])) > 1):
        fill_qty_ev_corr = \
            numpy.corrcoef(filldf['fill_qty'], filldf['next_fill_ret'])[0, 1]
      else:
        fill_qty_ev_corr = numpy.nan

      is_full_fill = filldf['fill_qty'] == filldf['fill_qty'].max()
      full_fill_ratio = is_full_fill.mean()
      fill_ev_bp = filldf['next_fill_ret'] / filldf['fill_price'] * 1e4

      agg_idx = ~numpy.isclose(filldf['fill_price'], filldf['price'], atol=0)

      for fwd in [30, 60]:
        profit = (filldf['next_fill_ret%s' % fwd] * filldf['fill_qty'])
        pq = (filldf['fill_price'] * filldf['fill_qty'])

        print("[%sF-fwd] %s agg_ratio: %s,"
              " agg_gross_ret_bp: %s,"
              " pass_gross_ret_bp: %s" %
              (fwd,
               oedf['symbol'][0],
               agg_idx.mean(),
               "{:,.2f}".format(profit[agg_idx].sum() / pq[agg_idx].sum() * 1e4),
               "{:,.2f}".format(profit[~agg_idx].sum() / pq[~agg_idx].sum() * 1e4)))

      if need_plot:
        matplotlib.rcParams.update({'font.size': 7})
        plt.subplot(311)
        plt.plot(pandas.DatetimeIndex(filldf['timestamp']), filldf['fill_price'], lw=0.5)
        plt.title('fill price, %s' % oedf['symbol'][0])
        plt.xticks(rotation=10)
        plt.subplot(312)

        plt.plot(pandas.DatetimeIndex(filldf['timestamp']),
                 (filldf['last_pos'].abs() * filldf['next_fill_ret']).cumsum(),
                 lw=1)
        plt.xticks(rotation=10)
        plt.twinx()
        plt.plot(pandas.DatetimeIndex(oedf['timestamp']), oedf['last_pos'], color='red', lw=0.5)
        plt.title('estimated pnl, position(red)')
        plt.subplot(325)
        plt.plot((filldf['fill_qty'] * filldf['fill_price']).cumsum())
        plt.title('cum turnover')
        plt.subplot(326)
        canqty = oedf['qty'] * (oedf['event_type'] == OrderEvent.CANCEL_CONFIRMED)
        plt.plot(pandas.DatetimeIndex(oedf['timestamp']), oedf['fill_qty'].cumsum(), 'g')
        plt.xticks(rotation=20)
        plt.twinx()
        plt.plot(pandas.DatetimeIndex(oedf['timestamp']), canqty.cumsum(), 'r')
        plt.title('fillqty/canqty sum')
        os.makedirs("oe_stat", exist_ok=True)
        tdate_str = pandas.DatetimeIndex(oedf['timestamp'])[0].strftime("%Y%m%d")
        plt.savefig("oe_stat/%s_%s.png" % (oedf['symbol'][0], tdate_str))
        plt.close()

        plt.subplot(211)
        plt.plot(pandas.DatetimeIndex(filldf['timestamp']),
                 filldf['fill_price'],
                 'g',
                 pandas.DatetimeIndex(filldf['timestamp']),
                 filldf['price'],
                 'r')
        plt.subplot(212)
        plt.plot(pandas.DatetimeIndex(filldf['timestamp'][agg_idx]),
                 filldf['fill_price'][agg_idx],
                 'r.')
        plt.twinx()
        plt.plot(pandas.DatetimeIndex(filldf['timestamp']), (filldf['fill_qty'] *
                                                             (filldf['fill_price'] > 0)).cumsum(),
                 'g')
        plt.plot(pandas.DatetimeIndex(filldf['timestamp']), (filldf['fill_qty'] *
                                                             (agg_idx)).cumsum(),
                 'g')
        plt.twinx()
        plt.plot(pandas.DatetimeIndex(filldf['timestamp']),
                 (fill_ev_bp * filldf['fill_qty']).cumsum(),
                 'y')
        plt.savefig("oe_stat/agg_view_%s_%s.png" % (oedf['symbol'][0], self._trading_date))
        plt.close()

      pnl_from_pos = (filldf['next_fill_ret'] * filldf['last_pos'].abs()).sum()
      fill_pq = (filldf['fill_qty'] * filldf['fill_price']).sum()

      fill_stats = {
          'agg_ratio_cnt':
              agg_idx.mean(),
          'agg_ratio_qty':
              0 if filldf['fill_qty'].sum() == 0 else filldf['fill_qty'][agg_idx].sum()
              / filldf['fill_qty'].sum(),
          'pnl_fill_est_from_pos':
              pnl_from_pos,
          'pnl_fill_est_from_fill':
              (filldf['last_fill_diff'] * filldf['fill_qty'] * filldf['sign']).sum() + first_pos *
              (prc_diff or 0),
          'full_fill_ratio':
              full_fill_ratio,
          'fill_qty':
              filldf['fill_qty'].sum(),
          'fill_qty_avg':
              filldf['fill_qty'].mean(),
          'fill_pq':
              fill_pq,
          'fill_qty_ratio':
              filldf['fill_qty'].sum() / subdf['qty'].sum(),
      }

      pnl_fee = get_pnl_fee(maker_pq=(filldf['fill_price'] * filldf['fill_qty'])[~agg_idx].sum(),
                            taker_pq=(filldf['fill_price'] * filldf['fill_qty'])[agg_idx].sum(),
                            market_type=self._market_type,
                            exchange=self._exchange)
      if pnl_fee is not None:
        fill_stats['pnl_fee'] = -pnl_fee
        fill_stats['pnl_fill_est_from_pos'] -= pnl_fee
        fill_stats['pnl_fill_est_from_fill'] -= pnl_fee

      if self._market_type == 'Spot':
        fill_stats['pnl_pos_reserve'] = pnl_pos_reserve

      fill_stats['pnl_fill_est_from_pos_ev_bp'] = (fill_stats['pnl_fill_est_from_pos']
                                                   / fill_stats['fill_pq'] * 1e4)

      if FLAGS.detail:
        fill_stats.update({
            'fill_qty_ev_corr':
                fill_qty_ev_corr,
            'fill_ev_bp_buy':
                get_weighted_mean_robust(fill_ev_bp[filldf['sign'] == 1],
                                         weights=filldf['fill_qty'][filldf['sign'] == 1]),
            'fill_ev_bp_sell':
                get_weighted_mean_robust(fill_ev_bp[filldf['sign'] == -1],
                                         weights=filldf['fill_qty'][filldf['sign'] == -1]),
            'fill_ev_partial_bp':
                get_weighted_mean_robust(fill_ev_bp[~is_full_fill],
                                         weights=filldf['fill_qty'][~is_full_fill]),
            'fill_ev_full_bp':
                get_weighted_mean_robust(fill_ev_bp[is_full_fill],
                                         weights=filldf['fill_qty'][is_full_fill]),
            'fill_qty_buy':
                filldf['fill_qty'][filldf['sign'] == 1].sum(),
            'fill_qty_sell':
                filldf['fill_qty'][filldf['sign'] == -1].sum(),
            'fill_qty_avg/lot_size_avg':
                filldf['fill_qty'].mean() / subdf['qty'].mean(),
        })
      statcol.update(fill_stats)

    return statcol

  def on_book_reset(self, book_builder_name, book_builder):
    for product in self._products:
      subscription_key = feed.get_subscription_key(product)
      if book_builder_name == subscription_key:
        book_builder.subscribe(product, functools.partial(self.on_book, product))
        if hasattr(book_builder, "subscribe_trade"):
          book_builder.subscribe_trade(product, functools.partial(self.on_trade, product))

  def try_assign_value(self, product, timestamp, col, val):
    symbol = self._product_map[product]
    if symbol not in self._stat_df:
      return
    sdf = self._stat_df[symbol]
    if sdf['timestamp'][0] > timestamp:
      return
    ind = sdf['timestamp'].searchsorted(timestamp, side='left')[0]
    if ind >= len(sdf):
      return
    if col not in self._arraydict[symbol]:
      self._arraydict[symbol][col] = numpy.array(sdf["timestamp"] * numpy.nan)
    self._arraydict[symbol][col][ind] = val

  def on_book(self, product, book):
    symbol = self._product_map[product]
    if symbol not in self._stat_df:
      return
    if self._last_book_timestamp + 5e8 > book.timestamp:
      return
    self._last_book_timestamp = book.timestamp
    exch = product.exchange
    self.try_assign_value(product, book.timestamp, "%s_ask0p" % exch, book.ask0().price)
    self.try_assign_value(product, book.timestamp, "%s_bid0p" % exch, book.bid0().price)

  def on_trade(self, product, trade):
    symbol = self._product_map[product]
    if symbol not in self._stat_df:
      return
    if self._last_trade_timestamp + 5e8 > trade.timestamp:
      return
    self._last_trade_timestamp = trade.timestamp
    exch = product.exchange
    if trade.side == TradeSide.TRADE_BUY_SIDE:
      self.try_assign_value(product, trade.timestamp, "%s_buyp" % exch, trade.price)
    elif trade.side == TradeSide.TRADE_SELL_SIDE:
      self.try_assign_value(product, trade.timestamp, "%s_sellp" % exch, trade.price)
    else:
      raise ValueError(trade.side)

  def on_log2(self, timestamp, log):
    pb = OrderGatewayLog()
    try:
      pb.ParseFromString(log)
    except Exception:
      print(log)
      return
    log_json = json_format.MessageToDict(pb, preserving_proto_field_name=True)
    self.on_log_json(timestamp, log_json)

  def update_currency_map(self, symbol):
    if symbol not in self._currency_map:
      if self._market_type in ['Spot', 'Margin']:
        self._currency_map[symbol] = SharedProduct.FromStr(symbol).base.currency
      elif self._market_type == 'Futures':
        if self._exchange == 'Okex':
          self._currency_map[symbol] = OkexFuturesProduct.FromStr(symbol,
                                                                  self._trading_date).base.currency
        elif self._exchange == 'Bitmex':
          self._currency_map[symbol] = BitmexFuturesProduct.FromStr(symbol).base.currency
        elif self._exchange == 'Bitflyer':
          self._currency_map[symbol] = BitflyerFuturesProduct.FromStr(symbol).base.currency
        else:
          raise ValueError(self._exchange)
      else:
        raise ValueError(self._market_type)

  def on_log_json(self, timestamp, log_json):
    if log_json['type'] in ['BALANCE', 'POSITION']:
      if log_json['type'] == 'BALANCE':
        if self._market_type in 'Spot' or self._exchange == 'Bitflyer':
          balance = log_json['balance']
          if 'each_balance' not in balance:
            return
          for each_balance in balance['each_balance']:
            if (self._market_type in 'Spot'
                or (self._exchange == 'Bitflyer' and each_balance['currency'] in ['JPY', 'BTC'])):
              self._balance_map[each_balance['currency']] = each_balance['total']
              self._balances[each_balance['currency']].append(
                  (timestamp, self._balance_map[each_balance['currency']]))
      elif log_json['type'] == 'POSITION' and self._market_type in ['Margin', 'Futures']:
        position = log_json['position']
        if 'each_position' in position:
          for each_position in position['each_position']:
            self.update_currency_map(each_position['symbol'])
            currency = self._currency_map[each_position['symbol']]
            self._balance_map[currency] = each_position['net_position']
            self._balances[currency].append((timestamp, self._balance_map[currency]))
    elif log_json['type'] == 'ORDER_EVENT':
      event = log_json['event']
      self._duration_tracker.on_log_json(timestamp, event)

      if self._exchange == 'Bitflyer':
        symbol = event['symbol']
        timestamp = int(event['event_time'])
        price = event.get('order_price', None)
        qty = event.get('order_qty', None)
        if 'internal_order_id' not in event:
          order_id = None
        else:
          order_id = self._order_id_converter.get_order_id_from_json(event)
      else:
        if 'internal_order_id' not in event:
          return
        symbol = event['symbol']
        timestamp = int(event['event_time'])
        if 'order_price' not in event:
          price = 0
        else:
          price = event['order_price']
        if 'order_qty' not in event:
          qty = 0
        else:
          qty = event['order_qty']
        order_id = self._order_id_converter.get_order_id_from_json(event)
        if order_id is None:
          return
      if 'order_side' not in event:
        return
      side = event['order_side']
      sign, side = convert_side(side)

      event_type = event['type']
      if event_type == 'ORDER_FILLED':
        event_type = OrderEvent.ORDER_FILLED
      elif event_type == 'ORDER_SUBMITTED':
        event_type = OrderEvent.ORDER_SUBMITTED
      elif event_type == 'CANCEL_CONFIRMED':
        event_type = OrderEvent.CANCEL_CONFIRMED
      elif event_type == 'ORDER_SUBMITTED':
        event_type = OrderEvent.ORDER_SUBMITTED

      fill_price = event['fill_price'] if 'fill_price' in event else 0
      fill_qty = event['fill_qty'] if 'fill_qty' in event else 0
      if price == 0:
        price = float('nan')
      if fill_price == 0:
        fill_price = float('nan')
      self.update_currency_map(symbol)
      base_currency = self._currency_map[symbol]
      last_pos = self._balance_map.get(base_currency, None)
      csv_line = [
          timestamp,
          symbol,
          event_type,
          price,
          qty,
          side,
          sign,
          order_id,
          fill_price,
          fill_qty,
          last_pos
      ]
      self._symbol_count[symbol] += 1
      self._rows[symbol].append(csv_line)
      if self._first_log_timestamp is None:
        self._first_log_timestamp = timestamp
      self._last_log_timestamp = timestamp

  def export_log(self):
    if self._exchange == 'Bitflyer':
      if len(self._balances['JPY']) > 0:
        ts, bals = zip(*self._balances['JPY'])
        ts_fut, bals_fut = zip(*self._balances['BTC'])
        print(pandas.DatetimeIndex([ts[-1]])[0])
        print("bitflyer est pnl: %s" % (bals[-1] - bals[0]))
        plt.plot(pandas.DatetimeIndex(ts), bals, lw=0.5)
        plt.title("%s jpy snapshot(blue)/ pos(red)" % self._machine)
        plt.twinx()
        plt.plot(pandas.DatetimeIndex(ts_fut), bals_fut, 'g', lw=0.2)
        plt.savefig("oe_stat/%s_%s_bitflyer_jpy.png" %
                    (self._machine, self._trading_date.strftime("%Y%m%d")))
        plt.close()
    for symbol in self._rows.keys():
      self._stat_df[symbol] = pandas.DataFrame(self._rows[symbol], columns=self._header)
      self._stat_oe[symbol] = self.get_stat_oe(self._stat_df[symbol],
                                               is_okex_futures=self._is_okex_futures)
      self._duration_tracker.report(symbol, need_plot=True)
      os.makedirs("oe_stat", exist_ok=True)
      self._stat_df[symbol].to_csv("oe_stat/%s_%s.csv" % (symbol, self._trading_date))
    print(
        pandas.DatetimeIndex([self._last_log_timestamp])[0],
        "\n",
        pandas.DataFrame(self._stat_oe).to_string())

  def print_futures_volume(self, file):
    assert self._exchange in ['Bitmex', 'Bitflyer']
    volume = {}
    for symbol in self._rows.keys():
      self._stat_df[symbol] = pandas.DataFrame(self._rows[symbol], columns=self._header)
      self._stat_oe[symbol] = self.get_stat_oe(self._stat_df[symbol],
                                               need_plot=False,
                                               is_okex_futures=self._is_okex_futures)
      volume[symbol] = {}
      try:
        if self._exchange == 'Bitmex':
          volume[symbol]['volume'] = self._stat_oe[symbol]['fill_qty']
        else:
          volume[symbol]['volume'] = self._stat_oe[symbol]['fill_pq']
      except KeyError:
        volume[symbol]['volume'] = 0

    formats = {}
    formats['volume'] = '{:.0f}'.format
    total_volume = 0
    for key in volume.keys():
      total_volume += volume[key]['volume']

    transposed = pandas.DataFrame(volume).T
    for key, value in formats.items():
      transposed[key] = transposed[key].apply(value)
    print(transposed.to_string(), file=file)

    volume_str = 'total volume: '
    if self._exchange == 'Bitmex':
      unit = 'USD'
    else:
      unit = 'JPY'
    volume_str += '%.0f %s' % (total_volume, unit)
    print(volume_str, file=file)
    print('first_log_time: %s' % datetime.datetime.fromtimestamp(self._first_log_timestamp / 1e9),
          file=file)
    print('last_log_time: %s' % datetime.datetime.fromtimestamp(self._last_log_timestamp / 1e9),
          file=file)

  def print_pnl_volume(self, file):
    pnl_volume = {}
    for symbol in self._rows.keys():
      self._stat_df[symbol] = pandas.DataFrame(self._rows[symbol], columns=self._header)
      self._stat_oe[symbol] = self.get_stat_oe(self._stat_df[symbol],
                                               need_plot=False,
                                               is_okex_futures=self._is_okex_futures)
      pnl_volume[symbol] = {}
      try:
        pnl_volume[symbol]['fill_pq'] = self._stat_oe[symbol]['fill_pq']
        pnl_volume[symbol]['pnl_fill_est'] = self._stat_oe[symbol]['pnl_fill_est_from_fill']
      except KeyError:
        pnl_volume[symbol]['fill_pq'] = 0
        pnl_volume[symbol]['pnl_fill_est'] = 0

    formats = {}
    total_volume = {}
    total_pnl = {}

    quote_unit = 'BTC'
    for key in pnl_volume.keys():
      unit = key.split('-')[1]
      total_volume.setdefault(unit, 0)
      total_volume[unit] += pnl_volume[key]['fill_pq']
      total_pnl.setdefault(unit, 0)
      total_pnl[unit] += pnl_volume[key]['pnl_fill_est']

    formats['agg_ratio_cnt'] = '{:.2f}'.format
    formats['agg_ratio_qty'] = '{:.2f}'.format
    formats['fill_qty_ratio'] = '{:.2f}'.format
    formats['full_fill_ratio'] = '{:.2f}'.format
    formats['pnl_fill_est_from_pos_ev_bp'] = '{:.4f}'.format
    if quote_unit == 'BTC':
      formats['fill_pq'] = '{:.3f}'.format
      formats['pnl_fill_est_from_fill'] = '{:.3f}'.format
      formats['pnl_fill_est_from_pos'] = '{:.3f}'.format
    else:
      formats['fill_pq'] = '{:.0f}'.format
      formats['pnl_fill_est_from_fill'] = '{:.0f}'.format
      formats['pnl_fill_est_from_pos'] = '{:.0f}'.format
    transposed = pandas.DataFrame(self._stat_oe).T
    for key, value in formats.items():
      transposed[key] = transposed[key].apply(value)
    print(transposed.to_string(), file=file)
    volume_str = 'total volume: '
    for unit, value in total_volume.items():
      if unit.upper() == 'BTC':
        volume_str += '%.3f %s    ' % (value, unit)
      else:
        volume_str += '%.0f %s    ' % (value, unit)
    print(volume_str, file=file)
    pnl_str = 'total pnl: '
    for unit, value in total_pnl.items():
      if unit.upper() == 'BTC':
        pnl_str += '%.3f %s    ' % (value, unit)
      else:
        pnl_str += '%.0f %s    ' % (value, unit)
    print(pnl_str, file=file)
    print('first_log_time: %s' % datetime.datetime.fromtimestamp(self._first_log_timestamp / 1e9),
          file=file)
    print('last_log_time: %s' % datetime.datetime.fromtimestamp(self._last_log_timestamp / 1e9),
          file=file)

  def export_log_feed(self):
    for symbol in self._arraydict.keys():
      sdf = pandas.DataFrame(self._stat_df[symbol])
      ndf = pandas.DataFrame(self._arraydict[symbol])
      tdf = pandas.concat([sdf, ndf], axis=1)
      tdf.to_csv("oe_stat/%s_%s_align.csv" % (symbol, self._trading_date))

  @property
  def products(self):
    if len(self._products) == 0:
      for symbol in self._symbol_count.keys():
        if self._market_type in ['Spot', 'Margin']:
          product = SharedProduct.FromStr(symbol)
        elif self._market_type == 'Futures':
          if self._exchange == 'Okex':
            product = OkexFuturesProduct.FromStr(symbol, self._trading_date)
          elif self._exchange == 'Bitmex':
            product = BitmexFuturesProduct.FromStr(symbol)
          elif self._exchange == 'Bitflyer':
            product = BitflyerFuturesProduct.FromStr(symbol)
          else:
            raise ValueError(self._exchange)

        products = []
        products.append(product)
        self._symbol_map[symbol] = products
        for product in products:
          self._product_map[product] = symbol
        self._products.extend(products)

    return self._products


def launch(exchange, trading_date, root_dir, machine):
  start_time = trading_date  # UTC 0 hour
  # 24 will omit some logs..
  hours = 25

  strategy = PtaStrategy(FLAGS.market_type, exchange, trading_date, machine)
  logging.info("start log analysis")
  run_from_og_log_archive(strategy.on_log2,
                          FLAGS.market_type,
                          exchange,
                          start_time,
                          start_time + datetime.timedelta(hours=hours),
                          root_dir,
                          machine,
                          owner=None)

  strategy.export_log()

  if FLAGS.align_with_feed:
    logging.info("start align with feed analysis")
    run_from_archive(strategy.products,
                     strategy.on_book_reset,
                     start_time,
                     start_time + datetime.timedelta(hours=hours),
                     machine=['feed-01.eu-west-1.aws'],
                     worker_id='1')

    strategy.export_log_feed()
  return 0


def get_trading_dates(trading_date_str):
  if trading_date_str.find("-") >= 0:
    tdfrom, tdto = trading_date_str.split("-")
    return pandas.date_range(tdfrom, tdto).to_pydatetime()
  else:
    return [datetime.datetime.strptime(trading_date_str, '%Y%m%d')]


def main(_):
  # import warnings; warnings.filterwarnings("error")
  exchange = FLAGS.exchange
  trading_date = FLAGS.trading_date
  root_dir = os.path.expanduser(FLAGS.root_dir)
  machine = FLAGS.machine
  assert FLAGS.market_type in ('ExWallet', 'Spot', 'Margin', 'Futures')
  assert trading_date, '--trading_date must be specified.'
  assert machine, '--machine must be specified.'
  assert exchange, '--exchange must be specified.'
  print('Running for %s %s %s ...' % (trading_date, machine, exchange))
  sys.stdout.flush()
  trading_dates = get_trading_dates(trading_date)
  if len(trading_dates) == 1:
    launch(exchange, trading_dates[0], root_dir, machine)
  else:
    # todo: multiprocessing
    for trading_date in trading_dates:
      launch(exchange, trading_date, root_dir, machine)


if __name__ == '__main__':
  flags.DEFINE_string('exchange', 'Okex', 'Exchange name.')

  flags.DEFINE_string('trading_date', None, 'Trading date in form of %Y%m%d.')

  flags.DEFINE_string('root_dir', '~/data/og_proto_log', 'root_dir.')

  flags.DEFINE_string('machine', 'strategy-04.aliyun-cn-hongkong', 'Instance machine name.')

  flags.DEFINE_string('market_type', 'Spot', 'Spot, Futures, Margin')

  flags.DEFINE_bool('align_with_feed', False, 'align with feed')

  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  app.run(main)
